from openai import OpenAI
import pymysql
from db import get_db
import re
import time
import hashlib
import json
import os
from datetime import datetime, timedelta

class RAGAssistant:
    """检索增强生成(Retrieval-Augmented Generation)助手"""
    
    def __init__(self, api_key, base_url="https://api.moonshot.cn/v1"):
        """初始化AI助手"""
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url
        )
        self.model = "moonshot-v1-8k"  # 默认模型
        
        # 请求限制相关
        self.last_request_time = 0
        self.min_request_interval = 20  # 每次请求至少间隔20秒
        
        # 缓存相关
        self.cache_dir = os.path.join(os.path.dirname(__file__), '..', 'cache')
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)
    
    def get_relevant_tables(self, query, limit=3):
        """根据查询获取相关的表名"""
        db = get_db()
        if not db:
            return []
            
        try:
            cursor = db.cursor()
            # 获取所有表名
            cursor.execute("SHOW TABLES")
            all_tables = [table[0] for table in cursor.fetchall()]
            cursor.close()
            
            # 简单关键词匹配查找相关表
            # 将查询拆分为关键词
            keywords = re.findall(r'\w+', query.lower())
            
            # 添加常见商品类别关键词映射 - 扩展更多类别和词汇
            keyword_mappings = {
                '家电': ['家电', 'appliance', '电器', '冰箱', '洗衣机', '空调', '电视', '微波炉', '电饭煲'],
                '家具': ['家具', 'furniture', '沙发', '桌子', '椅子', '床', '柜子', '衣柜', '书桌'], 
                '手机': ['手机', 'phone', '智能手机', '电话', 'iphone', '华为', '小米', '三星', 'oppo', 'vivo'],
                '电脑': ['电脑', 'computer', '笔记本', 'pc', 'laptop', '平板', '台式机', 'macbook', '联想', '戴尔'],
                '汽车': ['汽车', 'car', 'auto', '车', '轿车', 'suv', '二手车', '小车', '私家车', '代步'],
                '数码': ['数码', 'digital', '电子产品', '相机', '耳机', '音响', '智能手表', '无人机'],
                '办公': ['办公', 'office', '打印机', '复印机', '投影仪', '办公桌'],
                '母婴': ['母婴', '儿童', '婴儿', '玩具', '童车', '童装'],
                '运动': ['运动', '健身', '器材', '跑步机', '自行车', '球类'],
                '服装': ['服装', '衣服', '鞋子', '包包', '服饰', '男装', '女装'],
                '美妆': ['美妆', '化妆品', '护肤', '香水', '口红']
            }
            
            # 扩展关键词
            expanded_keywords = set(keywords)
            for keyword in keywords:
                for category, mappings in keyword_mappings.items():
                    if keyword in mappings:
                        expanded_keywords.add(category)
                        # 同时加入该类别的所有关键词，增强匹配能力
                        expanded_keywords.update(mappings)
            
            print(f"查询关键词: {keywords}")
            print(f"扩展后关键词: {expanded_keywords}")
            
            # 表名后缀和前缀模式
            table_patterns = ['_output', 'output_', 'out_put', '_data', 'data_', '_info', 'info_']
            
            # 按相关性对表进行评分和排序
            table_scores = []
            for table in all_tables:
                score = 0
                table_lower = table.lower()
                
                # 清理表名，移除常见后缀
                clean_table = table_lower
                for pattern in table_patterns:
                    clean_table = clean_table.replace(pattern, '')
                
                # 直接匹配完整表名（清理后）
                for keyword in expanded_keywords:
                    # 表名完全匹配某个关键词
                    if keyword == clean_table:
                        score += 10  # 完全匹配给更高分
                    # 表名包含关键词
                    elif keyword in clean_table:
                        score += 5  # 部分匹配也给分
                    # 表中字段匹配关键词
                    elif len(keyword) > 2 and any(keyword in field.lower() for field in self.get_table_columns(table)):
                        score += 2  # 字段匹配给较低分
                
                # 额外检查：如果原始查询中提到了某个类别，但表名是该类别的变种，给额外加分
                for category in keyword_mappings.keys():
                    if category in expanded_keywords and (
                        category in clean_table or 
                        any(pattern in table_lower for pattern in [f"{category}_", f"_{category}"])
                    ):
                        score += 8  # 类别相关性高的给额外加分
                
                if score > 0:
                    table_scores.append((table, score))
                    print(f"表 '{table}' 得分: {score}")
            
            # 按分数排序并返回前N个表
            relevant_tables = [t[0] for t in sorted(table_scores, key=lambda x: x[1], reverse=True)[:limit]]
            
            # 如果没有明显相关的表，返回一些通用的表
            if not relevant_tables and all_tables:
                # 尝试匹配通用商品表
                generic_tables = [t for t in all_tables if 'product' in t.lower() or '商品' in t.lower()]
                if generic_tables:
                    return generic_tables[:limit]
                # 否则返回前几个表
                return all_tables[:limit]
            
            print(f"匹配到的相关表: {relevant_tables}")
            return relevant_tables
        
        except Exception as e:
            print(f"获取相关表失败: {e}")
            return []
    
    def get_table_columns(self, table_name):
        """获取表的列名"""
        db = get_db()
        if not db:
            return []
            
        try:
            cursor = db.cursor()
            cursor.execute(f"DESCRIBE `{table_name}`")
            columns = [column[0] for column in cursor.fetchall()]
            cursor.close()
            return columns
        except Exception as e:
            print(f"获取表列名失败: {e}")
            return []
    
    def retrieve_table_data(self, table_name, limit=5):
        """检索表的示例数据"""
        db = get_db()
        if not db:
            return None
            
        try:
            cursor = db.cursor(pymysql.cursors.DictCursor)
            cursor.execute(f"SELECT * FROM `{table_name}` LIMIT {limit}")
            rows = cursor.fetchall()
            
            # 获取表结构
            schema = self.get_table_columns(table_name)
            
            cursor.close()
            
            return {
                "table_name": table_name,
                "schema": schema,
                "sample_data": rows
            }
        except Exception as e:
            print(f"检索表数据失败: {e}")
            return None
    
    def search_in_table(self, table_name, query, limit=7):
        """在特定表中搜索相关内容，使用扩展的关键词和智能过滤"""
        db = get_db()
        if not db:
            return []
        
        try:
            cursor = db.cursor(pymysql.cursors.DictCursor)
            
            # 获取表的列名
            columns = self.get_table_columns(table_name)
            
            # 扩展搜索词，同上面的关键词映射
            keywords = re.findall(r'\w+', query.lower())
            expanded_keywords = set(keywords)
            
            # 简化版的关键词映射，只包含最核心的词汇
            keyword_mappings = {
                '家电': ['电器', '冰箱', '洗衣机', '空调', '电视'],
                '家具': ['沙发', '桌子', '椅子', '床', '柜子'], 
                '手机': ['phone', '智能手机', 'iphone', '华为', '小米'],
                '电脑': ['computer', '笔记本', 'laptop', '台式机'],
                '汽车': ['car', 'auto', '轿车', 'suv', '二手车'],
                '数码': ['digital', '电子产品', '相机', '耳机']
            }
            
            # 扩展关键词
            for keyword in keywords:
                for category, mappings in keyword_mappings.items():
                    if keyword in mappings or keyword == category:
                        expanded_keywords.update(mappings)
                        expanded_keywords.add(category)
            
            # 构建更智能的搜索条件
            search_conditions = []
            params = []
            
            # 优先处理文本列
            text_columns = [col for col in columns if any(text_type in col.lower() 
                             for text_type in ['name', 'title', 'desc', 'brand', 'model', '名称', '标题', '描述', '品牌', '型号'])]
            
            # 如果有明确的文本列，优先在这些列中搜索
            if text_columns:
                for keyword in expanded_keywords:
                    if len(keyword) >= 2:  # 忽略太短的关键词
                        for col in text_columns:
                            search_conditions.append(f"`{col}` LIKE %s")
                            params.append(f"%{keyword}%")
            else:
                # 否则在所有列中搜索
                for keyword in expanded_keywords:
                    if len(keyword) >= 2:  # 忽略太短的关键词
                        for col in columns:
                            search_conditions.append(f"`{col}` LIKE %s")
                            params.append(f"%{keyword}%")
            
            # 如果没有搜索条件，返回随机记录
            if not search_conditions:
                search_query = f"SELECT * FROM `{table_name}` ORDER BY RAND() LIMIT {limit}"
                cursor.execute(search_query)
                results = cursor.fetchall()
                cursor.close()
                return results
            
            # 合并搜索条件
            search_query = f"SELECT * FROM `{table_name}` WHERE {' OR '.join(search_conditions)} LIMIT {limit}"
            
            # 执行搜索
            cursor.execute(search_query, params)
            results = cursor.fetchall()
            
            # 如果没有结果，返回随机数据作为示例
            if not results:
                cursor.execute(f"SELECT * FROM `{table_name}` ORDER BY RAND() LIMIT {limit}")
                results = cursor.fetchall()
            
            cursor.close()
            return results
            
        except Exception as e:
            print(f"表内搜索失败: {e}")
            return []
    
    def format_context_from_db(self, query):
        """从数据库构建上下文信息"""
        # 获取相关表
        relevant_tables = self.get_relevant_tables(query)
        
        if not relevant_tables:
            return "数据库中没有找到与查询相关的信息。"
        
        context_parts = []
        
        # 为每个相关表添加信息
        for table in relevant_tables:
            # 获取表数据
            table_data = self.retrieve_table_data(table)
            if not table_data:
                continue
                
            # 尝试在表中搜索更相关的内容
            search_results = self.search_in_table(table, query)
            
            # 构建此表的上下文
            table_context = f"### 表名: {table}\n"
            table_context += f"字段: {', '.join(table_data['schema'])}\n\n"
            
            # 添加搜索结果
            if search_results:
                table_context += "相关数据示例:\n"
                for i, row in enumerate(search_results[:3], 1):
                    row_str = ", ".join([f"{k}: {v}" for k, v in row.items()])
                    table_context += f"{i}. {row_str}\n"
            else:
                # 如果没有搜索结果，添加一些示例数据
                sample_data = table_data.get('sample_data', [])
                if sample_data:
                    table_context += "数据示例:\n"
                    for i, row in enumerate(sample_data[:2], 1):
                        row_str = ", ".join([f"{k}: {v}" for k, v in row.items()])
                        table_context += f"{i}. {row_str}\n"
            
            context_parts.append(table_context)
        
        # 组合所有上下文
        full_context = "\n\n".join(context_parts)
        return full_context
    
    def get_cache_key(self, query, context_type):
        """生成缓存键"""
        # 使用查询和上下文类型生成唯一键
        hash_input = f"{query}:{context_type}".encode('utf-8')
        return hashlib.md5(hash_input).hexdigest()
    
    def get_from_cache(self, cache_key):
        """从缓存获取回答"""
        cache_file = os.path.join(self.cache_dir, f"{cache_key}.json")
        
        if os.path.exists(cache_file):
            try:
                with open(cache_file, 'r', encoding='utf-8') as f:
                    cache_data = json.load(f)
                
                # 检查缓存是否过期（24小时）
                cache_time = datetime.fromisoformat(cache_data['timestamp'])
                if datetime.now() - cache_time < timedelta(hours=24):
                    print(f"从缓存获取回答: {cache_key[:8]}...")
                    return cache_data['response']
            except Exception as e:
                print(f"读取缓存错误: {e}")
        
        return None
    
    def save_to_cache(self, cache_key, response):
        """保存回答到缓存"""
        cache_file = os.path.join(self.cache_dir, f"{cache_key}.json")
        
        try:
            cache_data = {
                'timestamp': datetime.now().isoformat(),
                'response': response
            }
            
            with open(cache_file, 'w', encoding='utf-8') as f:
                json.dump(cache_data, f, ensure_ascii=False)
            
            print(f"回答已缓存: {cache_key[:8]}...")
        except Exception as e:
            print(f"保存缓存错误: {e}")
    
    def call_api_with_retry(self, messages, max_retries=3):
        """调用API并实现重试机制和速率限制"""
        # 检查距离上次请求的时间间隔
        current_time = time.time()
        time_since_last_request = current_time - self.last_request_time
        
        # 如果间隔太短，等待必要的时间
        if time_since_last_request < self.min_request_interval:
            wait_time = self.min_request_interval - time_since_last_request
            print(f"API速率限制: 等待 {wait_time:.1f} 秒...")
            time.sleep(wait_time)
        
        # 更新最后请求时间
        self.last_request_time = time.time()
        
        # 重试逻辑
        retries = 0
        last_error = None
        
        while retries < max_retries:
            try:
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    temperature=0.3,
                    max_tokens=800,
                )
                return response.choices[0].message.content
            
            except Exception as e:
                last_error = e
                retries += 1
                
                # 检查是否是速率限制错误
                if "rate_limit" in str(e).lower():
                    # 速率限制错误，等待更长时间
                    wait_time = (retries * 10) + 20  # 每次重试增加等待时间
                    print(f"API速率限制错误，等待 {wait_time} 秒后重试 ({retries}/{max_retries})...")
                    time.sleep(wait_time)
                else:
                    # 其他错误，短暂等待后重试
                    print(f"API错误: {e}, 等待 5 秒后重试 ({retries}/{max_retries})...")
                    time.sleep(5)
        
        # 所有重试都失败
        print(f"请求失败，已达最大重试次数: {last_error}")
        return f"抱歉，服务暂时繁忙，请稍后再试。(错误: {str(last_error)})"
    
    def get_response_with_preloaded_context(self, query, system_prompt):
        """使用预加载的数据库上下文获取AI回复（带缓存和重试）"""
        # 检查缓存
        cache_key = self.get_cache_key(query, "preloaded")
        cached_response = self.get_from_cache(cache_key)
        if cached_response:
            return cached_response
        
        # 首先判断查询是否与特定表相关
        relevant_tables = self.get_relevant_tables(query, limit=10)
        
        # 创建上下文和提示
        if relevant_tables:
            # 获取相关表的详细上下文
            db_context = self.format_context_from_db(query)
            enhanced_prompt = f"{system_prompt}\n\n以下是与问题最相关的数据库表信息：\n\n{db_context}"
            context_type = "relevant"
        else:
            # 获取预处理过的摘要版数据库信息
            db_summary = self.preprocess_all_tables(max_tables=10)  
            enhanced_prompt = f"{system_prompt}\n\n以下是数据库中最重要的表的信息：\n\n{db_summary}"
            context_type = "summary"
        
        # 构造消息
        messages = [
            {"role": "system", "content": enhanced_prompt},
            {"role": "user", "content": query}
        ]
        
        # 调用API (带重试)
        try:
            response = self.call_api_with_retry(messages)
            
            # 保存到缓存
            self.save_to_cache(cache_key, response)
            
            return response
            
        except Exception as e:
            print(f"AI响应错误: {e}")
            
            # 尝试回退到更简单的提示
            try:
                simplified_prompt = f"{system_prompt}\n\n请注意:数据库中有很多表，包括车辆、数码产品等二手商品信息。"
                messages = [
                    {"role": "system", "content": simplified_prompt},
                    {"role": "user", "content": query}
                ]
                
                fallback_response = self.call_api_with_retry(messages)
                return fallback_response
                
            except Exception as e2:
                return f"抱歉，服务暂时不可用，请稍后再试。(错误: {str(e2)})"

    def preprocess_all_tables(self, max_tables=10):
        """预处理有限数量的数据库表，生成摘要信息"""
        db = get_db()
        if not db:
            return "无法连接数据库"
        
        try:
            # 获取所有表名
            cursor = db.cursor()
            cursor.execute("SHOW TABLES")
            all_tables = [table[0] for table in cursor.fetchall()]
            
            # 获取表的行数，用于排序
            table_counts = []
            for table in all_tables:
                cursor.execute(f"SELECT COUNT(*) as count FROM `{table}`")
                count = cursor.fetchone()[0]
                table_counts.append((table, count))
            
            # 按记录数排序，选择最重要的表（记录数最多的）
            sorted_tables = sorted(table_counts, key=lambda x: x[1], reverse=True)
            selected_tables = [t[0] for t in sorted_tables[:max_tables]]
            
            cursor.close()
            
            # 存储所有表的信息
            tables_info = []
            
            # 添加表总览信息
            overview = f"数据库共有{len(all_tables)}个表，以下是最重要的{len(selected_tables)}个表的详细信息："
            tables_info.append(overview)
            
            for table in selected_tables:
                # 获取表的字段信息
                columns = self.get_table_columns(table)
                
                # 构建表信息 - 简化版，避免过多的示例数据
                table_info = f"表名: {table}\n"
                table_info += f"字段: {', '.join(columns)}\n"
                
                tables_info.append(table_info)
            
            # 组合所有表的信息
            all_tables_summary = "\n---\n".join(tables_info)
            
            return all_tables_summary
        
        except Exception as e:
            print(f"预处理数据库表错误: {e}")
            return f"在预处理数据库表时出错: {str(e)}"

    def preload_common_responses(self):
        """预加载常见问题的回答以提高响应速度"""
        common_questions = [
            "二手车怎么选", 
            "二手手机如何判断成色",
            "二手电子产品注意事项",
            "如何判断二手车好坏",
            "二手奢侈品鉴别方法",
            "如何评估二手车的价值",
            "二手电脑需要注意什么",
            "二手家电选购技巧"
        ]
        
        # 创建缓存目录
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)
        
        # 预加载常见问题的回答
        for question in common_questions:
            # 生成缓存键
            cache_key = self.get_cache_key(question, "preloaded")
            
            # 检查是否已缓存
            if self.get_from_cache(cache_key):
                print(f"已有缓存: {question}")
                continue
            
            # 没有缓存则生成回答并缓存
            try:
                print(f"预加载问题: {question}")
                
                # 为预加载回答添加图片示例
                system_prompt = """
                你是二手商品交易助手，提供关于二手车、数码产品等二手物品的专业建议。
                请在回答中添加1-2张相关的图片链接，格式为:
                ![描述](http://图片链接)
                """
                
                response = self.get_response_with_preloaded_context(
                    question, 
                    system_prompt
                )
                
                if response:
                    self.save_to_cache(cache_key, response)
                    print(f"成功预加载: {question}")
                
                # 避免API速率限制
                time.sleep(5)
                
            except Exception as e:
                print(f"预加载问题失败 '{question}': {e}")

    def process_question_with_db_examples(self, question):
        """基于用户问题匹配相关表并提取示例数据"""
        # 1. 获取相关表
        relevant_tables = self.get_relevant_tables(question, limit=2)
        
        if not relevant_tables:
            return None  # 没有匹配到相关表
        
        # 2. 从每个表中抽取示例数据
        db = get_db()
        if not db:
            return None
        
        examples_data = []
        try:
            cursor = db.cursor(pymysql.cursors.DictCursor)  # 使用字典游标获取列名
            
            for table in relevant_tables:
                # 获取表的列名
                columns = self.get_table_columns(table)
                
                # 随机抽取5条记录
                cursor.execute(f"SELECT * FROM `{table}` ORDER BY RAND() LIMIT 5")
                records = cursor.fetchall()
                
                if records:
                    # 格式化表数据
                    table_data = f"### 表 `{table}` 的示例数据:\n\n"
                    
                    # 添加表结构信息
                    table_data += f"表结构: {', '.join(columns)}\n\n"
                    
                    # 添加记录信息
                    for i, record in enumerate(records, 1):
                        table_data += f"**记录 {i}**:\n"
                        for key, value in record.items():
                            # 处理值，避免过长
                            if isinstance(value, str) and len(value) > 100:
                                value = value[:100] + "..."
                            table_data += f"- {key}: {value}\n"
                        table_data += "\n"
                    
                    examples_data.append(table_data)
            
            cursor.close()
            
        except Exception as e:
            print(f"获取表例子时出错: {e}")
            return None
        
        # 组合所有表的示例数据
        if examples_data:
            return "\n".join(examples_data)
        
        return None

    def get_response(self, query):
        """增强版获取回答函数，先检查缓存，没有再从LLM获取"""
        # 生成缓存键
        cache_key = self.get_cache_key(query)
        
        # 检查缓存
        cached_response = self.get_from_cache(cache_key)
        if cached_response and not query.lower().startswith("bypass cache:"):
            print(f"使用缓存回答: {cache_key[:8]}...")
            return cached_response
        
        # 如果请求包含 bypass cache 前缀，移除它
        if query.lower().startswith("bypass cache:"):
            query = query[len("bypass cache:"):].strip()
        
        # 没有缓存，需要生成新回答
        else:
            # 限制请求频率
            current_time = time.time()
            time_since_last_request = current_time - self.last_request_time
            
            if time_since_last_request < self.min_request_interval:
                sleep_time = self.min_request_interval - time_since_last_request
                print(f"等待 {sleep_time:.2f} 秒以符合API速率限制...")
                time.sleep(sleep_time)
            
            # 第一步：尝试从数据库获取示例数据
            db_examples = self.process_question_with_db_examples(query)
            
            # 设置system prompt
            system_prompt = """你是二手商品交易助手，提供关于二手车、数码产品等二手物品的专业建议。
            在回答中添加1-2张相关的图片链接，格式为:
            ![描述](图片URL)
            尽量基于提供的数据库信息回答，保持中立专业的语气。"""
            
            # 增强提示词
            if db_examples:
                enhanced_prompt = f"{system_prompt}\n\n以下是与您问题相关的数据库信息，请参考这些信息回答问题：\n\n{db_examples}"
                context_type = "relevant"
            else:
                # 获取预处理过的摘要版数据库信息
                db_summary = self.preprocess_all_tables(max_tables=10)  
                enhanced_prompt = f"{system_prompt}\n\n以下是数据库中最重要的表的信息：\n\n{db_summary}"
                context_type = "summary"
            
            # 构造消息
            messages = [
                {"role": "system", "content": enhanced_prompt},
                {"role": "user", "content": query}
            ]
            
            # 调用API (带重试)
            try:
                response = self.call_api_with_retry(messages)
                
                # 保存到缓存
                self.save_to_cache(cache_key, response)
                
                return response
                
            except Exception as e:
                print(f"AI响应错误: {e}")
                
                # 尝试回退到更简单的提示
                try:
                    simplified_prompt = f"{system_prompt}\n\n请注意:数据库中有很多表，包括车辆、数码产品等二手商品信息。"
                    messages = [
                        {"role": "system", "content": simplified_prompt},
                        {"role": "user", "content": query}
                    ]
                    
                    fallback_response = self.call_api_with_retry(messages)
                    return fallback_response
                    
                except Exception as e2:
                    return f"抱歉，服务暂时不可用，请稍后再试。"

    def clear_cache(self, pattern=None):
        """清除缓存文件
        
        参数:
            pattern: 可选的文件名匹配模式，如果为None则清除所有缓存
        """
        try:
            if not os.path.exists(self.cache_dir):
                print("缓存目录不存在，无需清理")
                return 0
            
            count = 0
            for file in os.listdir(self.cache_dir):
                if file.endswith('.json'):
                    if pattern is None or pattern in file:
                        os.remove(os.path.join(self.cache_dir, file))
                        count += 1
            
            print(f"成功清除 {count} 个缓存文件")
            return count
        except Exception as e:
            print(f"清除缓存时出错: {e}")
            return -1
